import tensorflow as tf
from tensorflow.keras import layers, models
import pandas as pd
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_curve
import matplotlib.pyplot as plt
import io

# 加载数据集
movies_df = pd.read_csv('movies.csv')

# 处理电影类型数据
genres = movies_df['genres'].apply(lambda x: x.split('|'))
mlb = MultiLabelBinarizer()
genres_encoded = mlb.fit_transform(genres)

# 将输入数据增加一个维度，以匹配卷积层期望的形状
genres_encoded = genres_encoded[:, :, None]

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(genres_encoded, genres_encoded[:, :, 0], test_size=0.2, random_state=42)

# 创建电影的特征输入
movie_input = tf.keras.Input(shape=(X_train.shape[1], X_train.shape[2]), name='Movie-Input')

# 创建网络层
# 一维卷积层
# 捕捉电影类别间的局部相关性。这个想法是将类别间的一些组合（如同时包含动作和科幻的电影）视作特征来捕捉。
conv1 = layers.Conv1D(filters=32, kernel_size=3, activation='relu')(movie_input)
pool1 = layers.MaxPooling1D(pool_size=2)(conv1)
conv2 = layers.Conv1D(filters=64, kernel_size=3, activation='relu')(pool1)
pool2 = layers.MaxPooling1D(pool_size=2)(conv2)
flat = layers.Flatten()(pool2)

# 全连接层
fc1 = layers.Dense(units=64, activation='relu')(flat)
output = layers.Dense(units=X_train.shape[1], activation='sigmoid')(fc1)

# 创建模型
model = models.Model(inputs=movie_input, outputs=output)

# 编译模型，加入AUC作为一个指标
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy', tf.keras.metrics.AUC(name='auc')])

# 准备TensorBoard的回调
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs4', histogram_freq=1, update_freq='batch')

# 训练模型
model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=10, callbacks=[tensorboard_callback])


import numpy as np
# 推荐函数
def recommend(movie_id, top_n=5):
    # 我们需要使用np.expand_dims来增加输入数据的维度
    movie_vec = genres_encoded[movie_id]
    movie_vec = np.expand_dims(movie_vec, axis=0)  # 这将数据的形状从(18,)转变为(1, 18)
    prediction = model.predict(movie_vec)
    prediction = np.squeeze(prediction)  # 把结果的维度从(1, 18)降低回(18,)

    # 获取排序后的最高评分电影的索引
    sorted_indexes = prediction.argsort()[::-1][:top_n]

    # 返回对应的电影标题
    recommended_movie_ids = movies_df.iloc[sorted_indexes]['movieId']
    return movies_df[movies_df['movieId'].isin(recommended_movie_ids)]['title']


# import pdb
# pdb.set_trace()
# 推荐与movieId为1的电影相似的5部电影
print(recommend(1))